import pandas as pd
import numpy as np

from torch.utils.data import Dataset


class arrhythmia_dataset_dev(Dataset):
    def __init__(self,train=True):
        if train :
            self.data = np.array(pd.read_csv('./dev/train.csv').iloc[:,np.arange(56)])
            label_not_tran = pd.read_csv('./dev/train.csv')['label']
            self.label = np.array(label_not_tran.map({0: 0, 1: 1}))
        else:
            self.data = np.array(pd.read_csv('./dev/test.csv').iloc[:, np.arange(56)])
            label_not_tran = pd.read_csv('./dev/test.csv')['label']
            self.label = np.array(label_not_tran.map({0: 0, 1: 1}))

    def __getitem__(self, index):
        return self.data[index], self.label[index]

    def __len__(self):
        return len(self.label)


class arrhythmia_dataset_equal(Dataset):
    def __init__(self,train=True):
        if train :
            self.data = np.array(pd.read_csv('./equal/train.csv').iloc[:, np.arange(56)])
            label_not_tran = pd.read_csv('./equal/train.csv')['label']
            self.label = np.array(label_not_tran.map({0: 0, 1: 1}))
        else:
            self.data = np.array(pd.read_csv('./equal/test.csv').iloc[:, np.arange(56)])
            label_not_tran = pd.read_csv('./equal/test.csv')['label']
            self.label = np.array(label_not_tran.map({0: 0, 1: 1}))

    def __getitem__(self, index):
        return self.data[index], self.label[index]

    def __len__(self):
        return len(self.label)

